import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision import transforms
from torchvision.transforms import ToTensor
import torchvision.transforms as tt
import numpy as np
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn import metrics

import os

# 导入自己创建的python文件
import sys
sys.path.append("..") # Adds higher directory to python modules path.
from frame.ModelUtil import NeuralNetwork, create_NN_model
from frame.DataProcess import CustomDataset, load_Purchase100
from frame.TrainUtil import train, train_models


# 这批参数用于resnet18
LEARNING_RATE = 1e-2
BATCH_SIZE = 128
MODEL = 'ResNet18'
EPOCHS = 100
DATA_NAME = 'CIFAR10'

# LEARNING_RATE = 1e-3
# BATCH_SIZE = 128
# MODEL = 'CNN'
# EPOCHS = 50
# DATA_NAME = 'CIFAR10'


# 批量训练模型
weight_dir = os.path.join('..', 'weights_for_exp', DATA_NAME)
# transform = transforms.Compose([])
# transform = transforms.Compose([
#     transforms.ToTensor(),
#     transforms.Normalize(mean=[0.507, 0.487, 0.441], std=[0.267, 0.256, 0.276])
#     ])
transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.RandomCrop(32, padding=4),  #先四周填充0，在吧图像随机裁剪成32*32
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.507, 0.487, 0.441], std=[0.267, 0.256, 0.276])
    ])
# transform = transforms.Compose([
#     transforms.ToTensor(),
#     transforms.Normalize(mean=[0.5], std=[0.5])
#     ])
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")
train_models(weight_dir, DATA_NAME, MODEL, LEARNING_RATE, EPOCHS, transform, device, num_shadowsets=100, prop_keep=0.5)